#!/bin/env python3
'''
Use this code directly:
    python write_fso_db_dev.py -ip 10.36.6.76 -t 2022022700 -i ../run/2022022700/obsimpact/gts_omb_oma_01

Note: pandas, psycopy2, sqlalchemy is the latest verion...

History:

    2022-3-3 (Rez)
    First commit

'''
import os
import re
import sys
from io import StringIO
from urllib import parse
import argparse
import pandas as pd
import psycopg2
import psycopg2.extras
from sqlalchemy import create_engine, text
import pendulum
import concurrent.futures
import time
import string
import requests

'''
DataFrame
0            1      2       3    4      5     6     7     8     9
'obs_type','var','stnid','lon','lat','pres','inv','obs','inc','time'
gts_omb_oma_01
0        1       2    3    4     5      6     7    8        9    
number   stnid   lat  lon  pres  u_obs  u_inv u_qc u_error  u_inc
                                 10     11    12   13       14
                                 v_obs  v_inv v_qc v_error  v_inc
                                 15     16    17   18       19
                                 t_obs  t_inv t_qc t_error  t_inc
                                 20     21    22   23       24
                                 p_obs  p_inv p_qc p_error  p_inc
                                 25     26    27   28       29
                                 q_obs  q_inv q_qc q_error  q_inc
'''
raw_header  = {}
raw_dtype   = {}
raw_header['synop'] = {'number':'i', 'stnid':str, 'lat':'f', 'lon':'f', 'pres':'f',
                       'u_obs':'f', 'u_inv':'f', 'u_qc':'i', 'u_error':'f', 'u_inc':'f',
                       'v_obs':'f', 'v_inv':'f', 'v_qc':'i', 'v_error':'f', 'v_inc':'f',
                       't_obs':'f', 't_inv':'f', 't_qc':'i', 't_error':'f', 't_inc':'f',
                       'p_obs':'f', 'p_inv':'f', 'p_qc':'i', 'p_error':'f', 'p_inc':'f',
                       'q_obs':'f', 'q_inv':'f', 'q_qc':'i', 'q_error':'f', 'q_inc':'f',}

raw_header['buoy']  = raw_header['synop']

raw_header['sound'] = {'number':int, 'stnid':str, 'lat':float, 'lon':float, 'pres':float,
                       'u_obs':float, 'u_inv':float, 'u_qc':int, 'u_error':float, 'u_inc':float,
                       'v_obs':float, 'v_inv':float, 'v_qc':int, 'v_error':float, 'v_inc':float,
                       't_obs':float, 't_inv':float, 't_qc':int, 't_error':float, 't_inc':float,
                       'q_obs':float, 'q_inv':float, 'q_qc':int, 'q_error':float, 'q_inc':float,}


raw_header['airep'] = raw_header['sound']

raw_header['profiler'] = {'number':int, 'stnid':str, 'lat':float, 'lon':float, 'pres':float,
                       'u_obs':float, 'u_inv':float, 'u_qc':int, 'u_error':float, 'u_inc':float,
                       'v_obs':float, 'v_inv':float, 'v_qc':int, 'v_error':float, 'v_inc':float,}


prod_header         = ['obs_type','var','stnid','lon','lat','lev','pres','inv','obs','qc','error','inc']



def generate_fake_stn_id(num):
    # 1.3 include duplicates member
    # a = [ ''.join(random.choices(string.ascii_uppercase + string.digits, k=5)) for _ in range(int(num * 1.3)) ]
    # remove duplicates
    # a = list(set(a))
    a = []
    cdict = string.digits + string.ascii_uppercase
    base  = len(cdict)
    for i in range(1,num+1):
        scode = []
        while True:
            idiv, imod = divmod(i, base)
            scode.insert(0, cdict[imod])
            if idiv == 0:
                break
            i = idiv
        a += [ f'{"".join(scode):>05s}' ]
    return a

def convert_table_format(obs_type, i_df):
    var_header = {'u':['stnid','lon','lat','pres','lev','u_inv','u_obs','u_qc','u_error','u_inc'],
                  'v':['stnid','lon','lat','pres','lev','v_inv','v_obs','v_qc','v_error','v_inc'],
                  't':['stnid','lon','lat','pres','lev','t_inv','t_obs','t_qc','t_error','t_inc'],
                  'q':['stnid','lon','lat','pres','lev','q_inv','q_obs','q_qc','q_error','q_inc'],
                  'p':['stnid','lon','lat','pres','lev','p_inv','p_obs','p_qc','p_error','p_inc'],}
    map_col    = {'u':{'u_inv':'inv','u_obs':'obs','u_qc':'qc','u_error':'error','u_inc':'inc'},
                  'v':{'v_inv':'inv','v_obs':'obs','v_qc':'qc','v_error':'error','v_inc':'inc'},
                  't':{'t_inv':'inv','t_obs':'obs','t_qc':'qc','t_error':'error','t_inc':'inc'},
                  'q':{'q_inv':'inv','q_obs':'obs','q_qc':'qc','q_error':'error','q_inc':'inc'},
                  'p':{'p_inv':'inv','p_obs':'obs','p_qc':'qc','p_error':'error','p_inc':'inc'},}
    o_df = pd.DataFrame(columns=prod_header)
    if obs_type == 'synop' or obs_type == 'buoy':
        for i in ['u','v','t','p','q']:
            tmp_df = i_df[var_header[i]].rename(columns=map_col[i])
            tmp_df['obs_type'] = obs_type
            tmp_df['var']      = i
            o_df = o_df.append(tmp_df)
    if obs_type == 'sound' or obs_type == 'airep':
        for i in ['u','v','t','q']:
            tmp_df = i_df[var_header[i]].rename(columns=map_col[i])
            tmp_df['obs_type'] = obs_type
            tmp_df['var']      = i
            o_df = o_df.append(tmp_df)
    if obs_type == 'profiler':
        for i in ['u','v']:
            tmp_df = i_df[var_header[i]].rename(columns=map_col[i])
            tmp_df['obs_type'] = obs_type
            tmp_df['var']      = i
            o_df = o_df.append(tmp_df)
    
    o_df = o_df[o_df['obs'] != -888888]
    if obs_type == 'airep':
        o_df = o_df[o_df['qc'] == 0]
    return o_df


def read_fso_impact(fn):
    
    res_df = pd.DataFrame(columns=prod_header)
    with open(fn, "r") as f:
        lines = f.read().splitlines()
        len_lines = len(lines)
        hand_loc = 0
        while True:
            obs_type, n_obs = lines[hand_loc].split()
            n_obs           = int(n_obs)
            if obs_type == 'synop' or obs_type == 'buoy':
                print(f'[Notice]: read {obs_type}, Total station number: {n_obs}, Now row number: {hand_loc}')
                raw_list    = lines[hand_loc+2:hand_loc+n_obs*2+1:2]
                raw_list    = [ i for i in raw_list if '\x00' not in i ]
                df          = pd.read_csv(StringIO('\n'.join(raw_list)), names=raw_header['synop'].keys(), dtype=raw_header['synop'], delim_whitespace=True)
                df['lev']   = 1
                df['stnid'] = df['stnid'].apply(lambda x: x[-5:])
                res_df      = res_df.append(convert_table_format(obs_type,df))
                #res_df.to_csv('2022030100.txt',float_format='%15.3f',index=False)
                hand_loc   += n_obs*2 + 1
            elif obs_type == 'sound':
                print(f'[Notice]: read {obs_type}, Total station number: {n_obs}, Now row number: {hand_loc}')
                for i in range(n_obs):
                    n_lev       = int(lines[hand_loc+1])
                    raw_list    = lines[hand_loc+2:hand_loc+2+n_lev]
                    if '\x00' in raw_list[0]:
                        hand_loc   += n_lev +1
                        continue
                    df          = pd.read_csv(StringIO('\n'.join(raw_list)), names=raw_header['sound'].keys(), dtype=raw_header['sound'], delim_whitespace=True)
                    df['lev']   = df['stnid'].apply(lambda x: x[:-5])
                    df['stnid'] = df['stnid'].apply(lambda x: x[-5:])
                    res_df      = res_df.append(convert_table_format(obs_type,df))
                    hand_loc   += n_lev +1
                hand_loc   += 1
            elif obs_type == 'airep':
                print(f'[Notice]: read {obs_type}, Total station number: {n_obs}, Now row number: {hand_loc}')
                fake_stnid  = generate_fake_stn_id(n_obs)
                for i in range(n_obs):
                    n_lev       = int(lines[hand_loc+1])
                    raw_list    = lines[hand_loc+2:hand_loc+2+n_lev]
                    if '\x00' in raw_list[0]:
                        hand_loc   += n_lev +1
                        continue
                    df          = pd.read_csv(StringIO('\n'.join(raw_list)), names=raw_header['sound'].keys(), dtype=raw_header['sound'], delim_whitespace=True)
                    df['lev']   = df['stnid'].apply(lambda x: x[:-5])
                    df['stnid'] = fake_stnid[i]
                    res_df      = res_df.append(convert_table_format(obs_type,df))
                    hand_loc   += n_lev +1
                hand_loc   += 1
            elif obs_type == 'profiler':
                print(f'[Notice]: read {obs_type}, Total station number: {n_obs}, Now row number: {hand_loc}')
                for i in range(n_obs):
                    n_lev       = int(lines[hand_loc+1])
                    raw_list    = lines[hand_loc+2:hand_loc+2+n_lev]
                    if '\x00' in raw_list[0]:
                        hand_loc   += n_lev +1
                        continue
                    df          = pd.read_csv(StringIO('\n'.join(raw_list)), names=raw_header['profiler'].keys(), dtype=raw_header['profiler'], delim_whitespace=True)
                    df['lev']   = df['stnid'].apply(lambda x: x[:-5])
                    df['stnid'] = df['stnid'].apply(lambda x: x[-5:])
                    res_df      = res_df.append(convert_table_format(obs_type,df))
                    hand_loc   += n_lev +1
                hand_loc   += 1
            else:
                # skip unknown obs type, eg: sonde_sfc
                print(f'[Warning]: skip {obs_type}, Now row number: {hand_loc}')
                while True:
                    hand_loc   += 1
                    if hand_loc >= len_lines or len(lines[hand_loc].split()) == 2:break
            
            if hand_loc >= len_lines:break
    return res_df.reset_index()

def run(df_type, tbl_type):
    '''
    0            1      2       3    4      5     6     7     8     9
    'obs_type','var','stnid','lon','lat','pres','inv','obs','inc','time'
    '''
    
    #tbl_name = "tb_n_{}_{}".format(tbl_type, args.time.format('YYYYMM'))
    tbl_name = "tb_n_{}_{}".format(tbl_type, args.time[0:6])
    #tbl_data_name = "tb_n_{}data_{}".format(tbl_type, args.time.format('YYYYMM'))
    tbl_data_name = "tb_n_{}data_{}".format(tbl_type, args.time[0:6])

    # delete old data
    time_1  =  time.time()
    del_sql = text(f"DELETE FROM {tbl_name} WHERE time={df_type.loc[0,'time']}")
    db_con.execute(del_sql)
    del_sql = text(f"DELETE FROM {tbl_data_name} WHERE time={df_type.loc[0,'time']}")
    db_con.execute(del_sql)
    time_2  =  time.time()
    print(f'[Notice]: delete time: {time_2 - time_1}')

    # get unique station name and id
    df_key = df_type[['time','stnid','lon','lat']].copy()
    df_key = df_key.drop_duplicates()
    df_key.to_sql(tbl_name, con=db_con, if_exists='append',index=False)
    df_res = pd.read_sql('select * from {} where time={}'.format(tbl_name,df_type.loc[0,'time']), con=db_con)
    # stnid length is 8 in Database, remove spaces
    df_res['stnid'] = df_res['stnid'].str.strip()
    id_dic = df_res.set_index('stnid')['id'].to_dict()
    # mapping id 
    df_type['id'] = df_type['stnid'].apply(lambda x: id_dic[x])
    df_type[['id','var','lev','pres','obs','inv','qc','error','inc','time']].to_sql(tbl_data_name, con=db_con, if_exists='append',index=False)

if __name__ == '__main__':
    '''
    write fso impact to pg database
    '''
    parser = argparse.ArgumentParser(description='write fso result to pg database...')
    # parser.add_argument('-t', '--time', help='file time (YYYYMMDDHH[MM]).', type=parse_time, required=True)
    parser.add_argument('-t', '--time', help='file time (YYYYMMDDHH[MM]).', required=True)
    parser.add_argument('-i', '--input', help='FSO result file.', required=True)
    parser.add_argument('-ip', help='Database ip address.', required=True)
    args = parser.parse_args()

    # Read impact file
    type_dic = {'synop':'surface','buoy':'buoy','sound':'sound', 'airep':'airep','profiler':'profiler'}
    col_name = ['obs_type','var','stnid','lon','lat','pres','inv','obs','inc']
    df = read_fso_impact(args.input)
    if len(df) < 1:
        print(f'[Error]: Empty file {args.input}!')
        exit(1)
    #args.time  = pendulum.from_format(args.time, 'YYYYMMDDHH')
    df['time']  = args.time #.format('YYYYMMDDHH')
    # Open Database
    engine   = create_engine('postgresql+psycopg2://fso:%s@%s/fso' % (parse.quote_plus('fSO@2017'), args.ip), pool_recycle=3600)
    db_con   = engine.connect()

    for obs_type in df.drop_duplicates(subset=['obs_type'])['obs_type']:
        if obs_type in type_dic.keys():
            print(obs_type)
            df_type = df.query(f'obs_type == "{obs_type}"')
            df_type = df_type.reset_index()
            run(df_type, type_dic[obs_type])
    db_con.close()
    #print('[Notice]: call fso result calculate REST')
    #post_url = f'http://10.36.6.75:9000/obps-server/api/rest/dailydata/chart_time/{args.time.format("YYYYMMDDHH")}'
    #r = requests.get(post_url)
    #print(r.content)
